#!/usr/bin/env python
# coding: utf-8

import os
import sys
import shutil
import envoy
import re
from metlib.shell.fileutil import BASENAME, DIRNAME, CP
from metlib.datetime import TR, T

def new_TR(beg, end=None, tdelta='1h', include_end=True):
    if include_end:
        tr = TR(beg, end, tdelta)
        tr.append(T(end))
        return tr
    else:
        return TR(beg, end, tdelta)


class ObsDataFinder(object):
    def __init__(self, obsdata_root):
        self._obsdata_root = obsdata_root

    def find_gzips(self, begdt, enddt):
        """find all obsdata files between begdt and enddt
        @:param begdt: begin datestr(yyyymmddhh)
        @:param enddt: end datestr(yyyymmddhh)
        """
        dt_range = new_TR(begdt, enddt, tdelta='1h')
        gzip_files = []
        for dt in dt_range:

            gzip_file = os.path.join(self._obsdata_root, dt.strftime('%Y%m%d'),
                                     '{}.szmb_decoded.gz'.format(dt.strftime('%Y%m%d%H')))

            if os.path.exists(os.path.join(self._obsdata_root, 'wmo_{}.decoded.gz'.format(dt.strftime('%Y%m%d%H')))):

                wmo_file = os.path.join(self._obsdata_root, 'wmo_{}.decoded.gz'.format(dt.strftime('%Y%m%d%H')))
            else:
                wmo_file = os.path.join(self._obsdata_root, dt.strftime('%Y%m%d'),
                                        'wmo_{}.decoded.gz'.format(dt.strftime('%Y%m%d%H')))
            madis_file = os.path.join(self._obsdata_root, dt.strftime('%Y%m%d'),
                                      'madis_{}.decoded.gz'.format(dt.strftime('%Y%m%d%H')))
            if os.path.exists(gzip_file):
                gzip_files.append(gzip_file)
            elif os.path.exists(wmo_file):
                gzip_files.append(wmo_file)
            elif os.path.exists(madis_file):
                gzip_files.append(madis_file)
            else:
                continue
        return gzip_files

    @staticmethod
    def copy_unzip_rename_files(gzip_files, dest=None, little_points=[], func=None):
        """unzip all gzipped files
        """
        unzip_files = []

        if not os.path.exists(dest):
            os.makedirs(dest)
        for gzip_file in gzip_files:
            CP(gzip_file, dest)
            new_gzip = os.path.join(dest, BASENAME(gzip_file))
            unzip_file = new_gzip.strip(".gz")
            if os.path.exists(unzip_file):
                os.remove(unzip_file)
            res = envoy.run("gzip -d %s" % new_gzip, cwd=dest)
            pattern = re.compile(r'\d+')
            if res.status_code == 0:

                dtstr = pattern.search(BASENAME(unzip_file)).group()
                obs_filename = 'obs:{}-{}-{}_{}'.format(dtstr[0:4], dtstr[4:6], dtstr[6:8], dtstr[8:10])
                print(obs_filename)
                obs_file = os.path.join(
                    dest, obs_filename)
                shutil.move(unzip_file, obs_file)
                if func is not None:
                    func(dest, obs_filename, little_points)
        #       obs_files.append(obs_file)
            else:
                print('corrupt obsfile: %s' % gzip_file)
        return unzip_files

    @staticmethod
    def rename_files(unzip_files):
        """rename unzipped files to obs:YYYY-MM-DD:HH format"""
        pattern = re.compile(r'\d+')
        dest = DIRNAME(unzip_files[0])
        obs_files = []
        for unzip_file in unzip_files:
            dtstr = pattern.search(BASENAME(unzip_file)).group()
            obs_file = os.path.join(
                dest, 'obs:{}-{}-{}_{}'.format(dtstr[0:4], dtstr[4:6], dtstr[6:8], dtstr[8:10]))
            shutil.move(unzip_file, obs_file)
            obs_files.append(obs_file)
        return obs_files

    def batch_process(self, begdt, enddt, dest=None, little_points=[], func=None):
        """do batch works to make obsgrid-ready data
        @:param begdt: begin datestr(yyyymmddhh)
        @:param enddt: end datestr(yyyymmddhh)
        @:param dest: optional destination path
        """
        gzip_files = self.find_gzips(begdt, enddt)
        unzip_files = self.copy_unzip_rename_files(gzip_files, dest, little_points, func)
#        obs_files = self.rename_files(unzip_files)


if __name__ == '__main__':

    if len(sys.argv) == 5:
        begdt = sys.argv[1]
        enddt = sys.argv[2]
        dest = sys.argv[3]
        obs_root = sys.argv[4]
    elif len(sys.argv) == 4:
        begdt = sys.argv[1]
        enddt = sys.argv[2]
        dest = sys.argv[3]
    elif len(sys.argv) == 3:
        begdt = sys.argv[1]
        enddt = sys.argv[2]
    else:
        print(
            "Usage: %s begdt(YYYMMDDHH) enddt [dest_path] [obs_root_path]" % sys.argv[0])
        exit(1)
    obs_finder = ObsDataFinder(obs_root)
    obs_finder.batch_process(begdt, enddt, dest)
